import asyncio
import os
from datetime import datetime
import redis
from metagpt.actions import Action, UserRequirement
from metagpt.logs import logger
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.environment import Environment

from metagpt.const import MESSAGE_ROUTE_TO_ALL
# 加载环境变量
from dotenv import load_dotenv
from metagpt.utils.common import OutputParser

load_dotenv()
novel = Environment()
red_db = redis.Redis(host='localhost', port=6379, db=0)



def save_file(context:Message):
    title = red_db.get("title").decode('utf-8')
    directory = "./workspace/"+title
    if not os.path.exists(directory):
        os.makedirs(directory)
    path = directory+"/"+context.content.split("\n")[0] +".txt"
    with open(path, 'w', encoding='utf-8') as file:
        # 写入字符串到文件
        file.write(context.content)


content_prompt_template = """
        根据{msg}
        您现在是短篇小说领域经验丰富的小说作家。我们需要您根据每一章标题和故事概况创作出小说详细的故事情节。
        请按照以下要求提供该小说的具体内容：
        1. 输出必须严格符合指定语言。
        2. 按照故事设定,章节标题，故事概况生成生动的故事内容和人物对话。
        3. 要求人物描述丰富，情节饱满。
        4. 输出内容最前面需要有章节和标题，如：第一章：章节标题 \n xxxx。
        4. 不要有额外的空格或换行符。
        5. 给我生成小说第{num}章，字数1000左右
        """

directory_prompt_template = """
        请根据{msg}
        您现在是短篇小说领域经验丰富的小说作家。请根据{msg} 内容创作出小说的目录和章节内容概况。
        请按照以下要求提供该小说的具体目录和目录中的故事概况：
        1. 输出必须严格符合指定语言。
        2. 回答按JSON数组格式输出，如: ["第一章：章节标题": "故事概况",\n "第二章：章节标题": "故事概况",\n]。
        3. 目录应尽可能引人注目和充分，包括一级目录和本章故事概况。
        4. 不要有额外的空格或换行符。
        """

main_prompt_template="""
        您现在是短篇小说领域经验丰富的小说作家, 我们需要您根据给定的小说题材生成故事的基本结构。
        按照以下内容输出{msg}题材的小说基本结构：
        标题:小说的标题
        设置:"小说的情景设置细节，包括时间段、地点和所有相关背景信息"
        主角:"小说主角的名字、年龄、职业，以及他们的性格和动机、简要的描述"
        反派角色:"小说反派角色的名字、年龄、职业，以及他们的性格和动机、简要的描述"
        冲突:"小说故事的主要冲突，包括主角面临的问题和涉及的利害关系"
        主题:"小说中心主题，并说明如何在整个情节、角色和背景中展开"
        基调:"整体故事的基调，以及保持背景和人物的一致性和适当性的说明"
        节奏:"调节故事节奏以建立和释放紧张气氛，推进情节，创造戏剧效果的说明"
        其它:"任何额外的细节或对故事的要求，如特定的字数或题材限制"
        """


class ContentAction(Action):
    """
    撰写Action。
    """
    name: str = "contentWriter"
    CONTENT_PROMPT_TEMPLATE: str = content_prompt_template
    async def run(self, msg: str, num : int):
        """
        根据用户提供的主题撰写一篇小说架构
        """
        num_str = str(num)
        prompt = self.CONTENT_PROMPT_TEMPLATE.format(msg=msg,num=num_str)
        print('content_prompt='+prompt)
        rsp = await self._aask(prompt)
        return rsp

class DirectoryAction(Action):
    """
    撰写Action。
    """
    name: str = "directoryWriter"
    PROMPT_TEMPLATE: str = directory_prompt_template
    async def run(self, msg: str):
        """
        根据用户提供的主题撰写一篇小说架构
        """
        prompt = self.PROMPT_TEMPLATE.format(msg=msg)
        rsp = await self._aask(prompt)
        return rsp


class MainAction(Action):
    """
    Agent的撰写小说的Action。
    """
    name: str = "titleWriter"
    PROMPT_TEMPLATE: str = main_prompt_template

    async def run(self, msg: str):
        """
        根据用户提供的主题撰写一篇小说架构
        """
        prompt = self.PROMPT_TEMPLATE.format(msg=msg)
        rsp = await self._aask(prompt)
        return rsp

class NovelContent(Role):
    """
    内容撰写角色。
    """
    name: str = "novelContent"
    profile: str = "NovelContent"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([ContentAction])  # 设置动作
        self._watch([DirectoryAction])  # 监听用户要求

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: ready to {self.rc.todo}")
        todo = self.rc.todo
        msg = self.get_memories()  # 获取所有对话记忆
        arr = red_db.get("tiles").decode('utf-8').split(",\n")
        print("arr===="+str(arr))
        for i in range(0,len(arr)):
            essay_text = await ContentAction().run(msg, i+1)
            logger.info(f'NovelContent : {essay_text}')
            msg = Message(content=essay_text, role=self.profile, cause_by=type(todo))
            # self.put_message(msg) #由于费用问题，这里每一章的内容不加入msg，以免消耗过多的token
            save_file(msg)
        return msg




class NovelDirectory(Role):
    """
    生成目录角色。
    """
    name: str = "novelDirectory"
    profile: str = "NovelDirectory"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([DirectoryAction])  # 设置动作
        self._watch([MainAction])  # 监听用户要求

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: ready to {self.rc.todo}")
        todo = self.rc.todo
        msg = self.get_memories()  # 获取所有对话记忆
        essay_text = await DirectoryAction().run(msg)
        logger.info(f'NovelDirectory : {essay_text}')
        msg = Message(content=essay_text, role=self.profile, cause_by=type(todo))
        print("msg.content="+msg.content)
        red_db.set("tiles", msg.content)
        self.put_message(msg)
        return msg


class NovelMain(Role):
    """
    小说架构师角色。
    """
    name: str = "novelMain"
    profile: str = "NovelMain"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([MainAction])  # 设置动作为撰写作
        self._watch([UserRequirement])  # 监听用户要求

    async def _act(self) -> Message:
        """
        小说架构师动作：根据用户要求撰写小说主要信息
        """
        logger.info(f"{self._setting}: ready to {self.rc.todo}")
        todo = self.rc.todo

        msg = self.get_memories()  # 获取所有对话记忆
        # logger.info(msg)
        essay_text = await MainAction().run(msg)
        title = essay_text.split("\n")[0].replace("标题:","")
        red_db.set("title",title)
        logger.info(f'NovelMain : {essay_text}')
        msg = Message(content=essay_text, role=self.profile, cause_by=type(todo))
        self.put_message(msg)
        return msg


async def main(topic: str, n_round=5):
    """
    运行函数，用户输入一个主题，并将主题发布在环境中，然后运行环境。
    """
    novel.add_roles([NovelMain(), NovelDirectory(), NovelContent()])  # 向环境中添加agent

    novel.publish_message(
        Message(role="Human", content=topic, cause_by=UserRequirement,
                send_to='' or MESSAGE_ROUTE_TO_ALL),
        peekable=False,
    )
    # 发布一条消息，包含用户输入的主题，并将其发送给所有角色
    while n_round > 0:
        n_round -= 1
        logger.debug(f"max {n_round=} left.")  # 输出剩余对话轮数
        await novel.run()  # 运行环境
    return novel.history  # 返回对话历史记录




if __name__ == '__main__':

    asyncio.run(main(topic='古代修仙'))  # 运行主函数